#include "GridOperator.cuh"
#include "GridSampler.cuh"

// This kernel evaluates v = s1 * v1 + s2 on MAC grids
// The values of the scalar field are defined at cell centers
template<typename T>
__global__ void saddKernel(int nx, int ny, int nz, int z, T s1, T *v1, T s2, T *v)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int idx = x + y * nx + z * nx * ny;
		v[idx] = s1 * v1[idx] + s2;
	}
}

template<typename T>
void sadd(int nx, int ny, int nz, T s1, T *v1, T s2, T *v)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		saddKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, s1, v1, s2, v);
	cutilCheckMsg("addKernel() execution failed!");
}

template
void sadd<float>(int nx, int ny, int nz, float s1, float *v1, float s2, float *v);

// This kernel evaluates v = s1 * v1 + s2 * v2 on MAC grids
// The values of scalar fields are defined at cell centers
template<typename T>
__global__ void vaddKernel(int nx, int ny, int nz, int z, T s1, T *v1, T s2, T *v2, T *v)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int idx = x + y * nx + z * nx * ny;
		v[idx] = s1 * v1[idx] + s2 * v2[idx];
	}
}

template<typename T>
void vadd(int nx, int ny, int nz, T s1, T *v1, T s2, T *v2, T *v)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		vaddKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, s1, v1, s2, v2, v);
	cutilCheckMsg("addKernel() execution failed!");
}

template
void vadd<float>(int nx, int ny, int nz, float s1, float *v1, float s2, float *v2, float *v);

// This function evaluates v = s1 * v1 on MAC grids
// The values of scalar fields are defined at cell centers
template<typename T>
void smul(int nx, int ny, int nz, T s1, T *v1, T *v)
{
	sadd(nx, ny, nz, s1, v1, T(0), v);
}

template
void smul<float>(int nx, int ny, int nz, float s1, float *v1, float *v);

// This kernel evaluates v = s1 * v1 * v2 on MAC grids
// The values of scalar fields are defined at cell centers
template<typename T>
__global__ void vmulKernel(int nx, int ny, int nz, int z, T s1, T *v1, T *v2, T *v)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int idx = x + y * nx + z * nx * ny;
		v[idx] = s1 * v1[idx] * v2[idx];
	}
}

template<typename T>
void vmul(int nx, int ny, int nz, T s1, T *v1, T *v2, T *v)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		vmulKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, s1, v1, v2, v);
	cutilCheckMsg("addKernel() execution failed!");
}

template
void vmul<float>(int nx, int ny, int nz, float s1, float *v1, float *v2, float *v);

// This kernel evaluates the divergence of vector fields on MAC grids
// The values of vector fields are defined at face centers
template<typename T>
__global__ void divergenceKernel(int nx, int ny, int nz, int z, T dx, T *u, T *v, T *w, T *div)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int idx = x + y * nx + z * nx * ny;
		T ux = u[idx + y + z * ny + 1] - u[idx + y + z * ny];
		T vy = v[idx + z * nx + nx] - v[idx + z * nx];
		T wz = w[idx + nx * ny] - w[idx];
		div[idx] = (ux + vy + wz) / dx;
	}
}

template<typename T>
void divergence(int nx, int ny, int nz, T dx, T *u, T *v, T *w, T *div)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		divergenceKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, dx, u, v, w, div);
	cutilCheckMsg("divergenceKernel() execution failed!");
}

template
void divergence<float>(int nx, int ny, int nz, float dx, float *u, float *v, float *w, float *div);

// This kernel evaluates the curl of vector fields on MAC grids
// The values of vector fields are defined at cell centers
// cx = wy - vz
// cy = uz - wx
// cz = vx - uy
template<typename T>
__global__ void curlKernel(int nx, int ny, int nz, int z, T dx, T *u, T *v, T *w, T *cx, T *cy, T *cz)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int nxy = nx * ny;
		int idx = x + y * nx + z * nxy;
		int id1 = (x == 0 ? idx : idx - 1);
		int id2 = (x == nx - 1 ? idx : idx + 1);
		int id3 = (y == 0 ? idx : idx - nx);
		int id4 = (y == ny - 1 ? idx : idx + nx);
		int id5 = (z == 0 ? idx : idx - nxy);
		int id6 = (z == nz - 1 ? idx : idx + nxy);
		T uy = u[id4] - u[id3];
		T uz = u[id6] - u[id5];
		T vz = v[id6] - v[id5];
		T vx = v[id2] - v[id1];
		T wx = w[id2] - w[id1];
		T wy = w[id4] - w[id3];
		T r2dx = T(0.5) / dx;
		cx[idx] = (wy - vz) * r2dx;
		cy[idx] = (uz - wx) * r2dx;
		cz[idx] = (vx - uy) * r2dx;
	}
}

template<typename T>
void curl(int nx, int ny, int nz, T dx, T *u, T *v, T *w, T *cx, T *cy, T *cz)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		curlKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, dx, u, v, w, cx, cy, cz);
	cutilCheckMsg("curlKernel() execution failed!");
}

template
void curl<float>(int nx, int ny, int nz, float dx, float *u, float *v, float *w, float *cx, float *cy, float *cz);

// This kernel evaluates the L2-norms of vector fields on MAC grids
// The values of vector fields are defined at cell centers
template<typename T>
__global__ void l2normKernel(int nx, int ny, int nz, int z, T dx, T *u, T *v, T *w, T *n)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int idx = x + y * nx + z * nx * ny;
		T ui = u[idx];
		T vi = v[idx];
		T wi = w[idx];
		n[idx] = sqrt(ui * ui + vi * vi + wi * wi);
	}
}

template<typename T>
void l2norm(int nx, int ny, int nz, T dx, T *u, T *v, T *w, T *n)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		l2normKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, dx, u, v, w, n);
	cutilCheckMsg("l2normKernel() execution failed!");
}

template
void l2norm<float>(int nx, int ny, int nz, float dx, float *u, float *v, float *w, float *n);

// This kernel resamples vector fields at cell centers on MAC grids
// The original values of vector fields are defined at face centers
template<typename T>
__global__ void resampleKernel(int nx, int ny, int nz, int z, T dx, T *u, T *v, T *w, T *uout, T *vout, T *wout)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;
	if(x < nx && y < ny)
	{
		int idx = x + y * nx + z * nx * ny;
		uout[idx] = (u[idx + y + z * ny] + u[idx + y + z * ny + 1]) * T(0.5);
		vout[idx] = (v[idx + z * nx] + v[idx + z * nx + nx]) * T(0.5);
		wout[idx] = (w[idx] + w[idx + nx * ny]) * T(0.5);
	}
}

template<typename T>
void resample(int nx, int ny, int nz, T dx, T *u, T *v, T *w, T *uout, T *vout, T *wout)
{
	dim3 gdim((nx - 1) / BLOCK_DIM + 1, (ny - 1) / BLOCK_DIM + 1);
	dim3 bdim(BLOCK_DIM, BLOCK_DIM, 1);
	for(int z = 0; z < nz; ++z)
		resampleKernel<T><<<gdim, bdim>>>(nx, ny, nz, z, dx, u, v, w, uout, vout, wout);
	cutilCheckMsg("resampleKernel() execution failed!");
}

template
void resample<float>(int nx, int ny, int nz, float dx, float *u, float *v, float *w, float *uout, float *vout, float *wout);
/*
template<typename T, bool CX, bool CY, bool CZ>
__global__ void sampleParticlesKernel(int nx, int ny, int nz, T dx, T *f, int np, float3 *p, float4 *data)
{
	int idx = blockIdx.x * blockDim.x + threadIdx.x;
	if(idx < np)
	{
		float3 pi = p[idx];
		T rdx = T(1) / dx;
		T xi = pi.x * rdx;
		T yi = pi.y * rdx;
		T zi = pi.z * rdx;
		data[idx] = sample<T>(nx, ny, nz, CX ? xi - T(0.5) : xi, CY ? yi - T(0.5) : yi, CZ ? zi - T(0.5) : zi, f);
	}
}

template<typename T>
void sampleParticles(int nx, int ny, int nz, T dx, T *f, int np, float3 *p, float4 *data, const FieldType type)
{
	const int bdim2 = BLOCK_DIM * BLOCK_DIM;
	dim3 gdim((np - 1) / bdim2 + 1, 1);
	dim3 bdim(bdim2, 1, 1);
	switch(type)
	{
	case CellCenter:
		sampleParticlesKernel<T, true, true, true><<<gdim, bdim>>>(nx, ny, nz, dx, f, np, p, data);
		break;
	case FaceCenterX:
		sampleParticlesKernel<T, false, true, true><<<gdim, bdim>>>(nx, ny, nz, dx, f, np, p, data);
		break;
	case FaceCenterY:
		sampleParticlesKernel<T, true, false, true><<<gdim, bdim>>>(nx, ny, nz, dx, f, np, p, data);
		break;
	case FaceCenterZ:
		sampleParticlesKernel<T, true, true, false><<<gdim, bdim>>>(nx, ny, nz, dx, f, np, p, data);
		break;
	}
	cutilCheckMsg("sampleParticlesKernel() execution failed!");
}

template
void sampleParticles<float>(int nx, int ny, int nz, float dx, float *f, int np, float3 *p, float4 *data, const FieldType type);
*/